# -*- coding:utf-8 -*-  
# Pwn module: full exploit generation 

from ropgenerator.IO import string_bold, string_special, banner, error, verbose_mode, verbose, string_exploit, remove_colors
from ropgenerator.Constraints import Constraint, Assertion, BadBytes
import ropgenerator.Architecture as Arch
from ropgenerator.exploit.pwn.DeliverShellcode import dshell
import ropgenerator.Database as Database
from ropgenerator.semantic.Engine import getBaseAssertion
from ropgenerator.exploit.Utils import parse_bad_bytes
from ropgenerator.CommonUtils import set_offset, reset_offset
import sys

###################
#  PWN COMMAND   # 
###################

# Sub-commands
CMD_DELIVER_SHELLCODE = "deliver-shellcode"

# Options 
OPTION_OUTPUT = '--output-format'
OPTION_OUTPUT_SHORT = '-f'
# Types for output
OUTPUT_CONSOLE = 'console'
OUTPUT_PYTHON = 'python'
OUTPUT_RAW = 'raw'
OUTPUT = None # The one choosen 

OPTION_BAD_BYTES = '--bad-bytes'
OPTION_BAD_BYTES_SHORT = '-b'

OPTION_HELP = "--help"
OPTION_HELP_SHORT = "-h"

OPTION_LMAX = '--max-length'
OPTION_LMAX_SHORT = '-m'

OPTION_VERBOSE = "--verbose"
OPTION_VERBOSE_SHORT = "-v"

OPTION_OUTFILE = '--output-file'
OPTION_OUTFILE_SHORT = '-o'

OPTION_PADDING_BYTE = '--padding-byte'
OPTION_PADDING_BYTE_SHORT = '-pb'

OPTION_PADDING_LEN = '--padding-len'
OPTION_PADDING_LEN_SHORT = '-pl'

OPTION_OFFSET = '--offset'
OPTION_OFFSET_SHORT = '-off'

CMD_PWN_HELP =  banner([string_bold("'pwn' command"),\
                    string_special("(Generate full exploits)")])
CMD_PWN_HELP += "\n\n\t"+string_bold("Usage:")+\
"\n\t\tpwn [OPTIONS] <subcommand> [SUBCOMMAND_OPTIONS]"

CMD_PWN_HELP += "\n\n\t"+string_bold("Subcommands")+":"
CMD_PWN_HELP += "\n\t(For more info use 'pwn <subcommand> -h')"
CMD_PWN_HELP += "\n\n\t\t"+string_special(CMD_DELIVER_SHELLCODE)+"\t Inject a shellcode an execute it"

CMD_PWN_HELP += "\n\n\t"+string_bold("Options")+":"
CMD_PWN_HELP += "\n\n\t\t"+string_special(OPTION_BAD_BYTES_SHORT)+","+string_special(OPTION_BAD_BYTES)+" <bytes>\t Bad bytes for payload.\n\t\t\t\t\t Expected format is a list of bytes \n\t\t\t\t\t separated by comas (e.g '-b 0A,0B,2F')"
CMD_PWN_HELP += "\n\n\t\t"+string_special(OPTION_LMAX_SHORT)+","+string_special(OPTION_LMAX)+" <int>\t Max length of the ROPChain in bytes"
CMD_PWN_HELP += "\n\n\t\t"+string_special(OPTION_PADDING_BYTE_SHORT)+","+string_special(OPTION_PADDING_BYTE)+" <byte> Byte for payload padding"
CMD_PWN_HELP += "\n\n\t\t"+string_special(OPTION_PADDING_LEN_SHORT)+","+string_special(OPTION_PADDING_LEN)+" <int>\t Length of payload padding"
CMD_PWN_HELP += "\n\n\t\t"+string_special(OPTION_OFFSET_SHORT)+","+string_special(OPTION_OFFSET)+" <int>\t Offset to add to gadget addresses"
CMD_PWN_HELP += "\n\n\t\t"+string_special(OPTION_OUTPUT_SHORT)+","+\
    string_special(OPTION_OUTPUT)+\
    " <fmt> Output format for ropchains.\n\t\t\t\t\t Expected format is one of the\n\t\t\t\t\t following: "+\
    string_special(OUTPUT_CONSOLE)+','+string_special(OUTPUT_PYTHON)
CMD_PWN_HELP += "\n\n\t\t"+string_special(OPTION_OUTFILE_SHORT)+","+string_special(OPTION_OUTFILE)+" <file>\t Save payload in a file"
CMD_PWN_HELP += "\n\n\t\t"+string_special(OPTION_HELP_SHORT)+","+string_special(OPTION_HELP)+"\t\t Show this help"
CMD_PWN_HELP += "\n\n\t"+string_bold("Examples")+": "+\
    "\n\t\tNo examples yet"
    
def print_help():
    print(CMD_PWN_HELP)


########################################################################
EXPLOIT_LMAX = 2000

def pwn(args):
    global OUTPUT, OUTPUT_CONSOLE, OUTPUT_PYTHON
    global EXPLOIT_LMAX
    
    
    if( not args ):
        print_help()
        return 
        
    # Check if gadgets have been loaded in the DB 
    if( not Database.gadgets ):
        error("Oops. You have to load gadgets before pwning ;) ")
        return 
        
    # Set default output format 
    OUTPUT = OUTPUT_CONSOLE
    clmax = EXPLOIT_LMAX
    offset = 0 
    
    seenOffset = False
    seenOutput = False
    seenLmax = False
    seenBadBytes = False
    seenVerbose = False
    seenOutfile = False
    seenPaddingByte = False
    seenPaddingLen = False

    constraint = Constraint()
    assertion = getBaseAssertion()
    subcommand = None
    outfile = None
    paddingByteStr = None 
    paddingLen = 0
    
    i = 0 
    while i < len(args):
        # Parse options     
        if( args[i][0] == '-' ):
            if( args[i] in [OPTION_BAD_BYTES, OPTION_BAD_BYTES_SHORT]):
                if( seenBadBytes ):
                    error("Error. '" + args[i] + "' option should be used only once")
                    return 
                if( i+1 >= len(args)):
                    error("Error. Missing bad bytes after option '"+args[i]+"'")
                    return 
                seenBadBytes = True
                (success, res) = parse_bad_bytes(args[i+1])
                if( not success ):
                    error(res)
                    return
                i = i+2
                constraint = constraint.add(BadBytes(res))
            elif( args[i] == OPTION_LMAX or args[i] == OPTION_LMAX_SHORT ):
                if( seenLmax ):
                    error("Error. '" + arg + "' option should be used only once.")
                    return 
                if( i+1 >= len(args)):
                    error("Error. Missing output format after option '"+arg+"'")
                    return 
                try:
                    clmax = int(args[i+1])
                    if( clmax < Arch.octets() ):
                        raise Exception()
                    # Convert number of bytes into number of ropchain elements
                    clmax /= Arch.octets()
                except:
                    error("Error. '" + args[i+1] +"' bytes is not valid")
                    return 
                i = i +2
                seenLmax = True
                       
            elif( args[i] in [OPTION_OUTPUT, OPTION_OUTPUT_SHORT]):
                if( seenOutput ):
                    error("Option '{}' should be used only once".format(args[i]))
                    return 
                if( i+1 >= len(args)):
                    error("Error. Missing output format after option '"+args[i]+"'")
                    return 
                if( args[i+1] in [OUTPUT_CONSOLE, OUTPUT_PYTHON]):
                    OUTPUT = args[i+1]
                    seenOutput = True
                    i += 2
                else:
                    error("Error. Unknown output format: {}".format(args[i+1]))
                    return 
            elif( args[i] in [OPTION_OUTFILE_SHORT, OPTION_OUTFILE] ):
                if( seenOutfile ):
                    error("Option '{}' should be used only once".format(args[i]))
                    return 
                if( i+1 >= len(args)):
                    error("Error. Missing file name after option '"+args[i]+"'")
                    return 
                seenOutfile = True
                outfile = args[i+1]
                i += 2
            elif( args[i] in [OPTION_PADDING_BYTE, OPTION_PADDING_BYTE_SHORT] ):
                if( seenPaddingByte ):
                    error("Option '{}' should be used only once".format(args[i]))
                    return 
                if( i+1 >= len(args)):
                    error("Error. Missing byte after option '"+args[i]+"'")
                    return 
                seenPaddingByte = True
                try:
                    paddingByte = int(args[i+1], 16)
                except:
                    error("Error. '{}' is not a valid byte".format(args[i+1]))
                    return 
                if( paddingByte > 0xff ):
                    error("Error. '{}' can't be stored in one byte".format(args[i+1]))
                    return
                paddingByteStr = "\\x"+format(paddingByte, "02x")
                i += 2
            elif( args[i] in [OPTION_PADDING_LEN, OPTION_PADDING_LEN_SHORT] ):
                if( seenPaddingLen ):
                    error("Option '{}' should be used only once".format(args[i]))
                    return 
                if( i+1 >= len(args)):
                    error("Error. Missing length after option '"+args[i]+"'")
                    return 
                seenPaddingLen = True
                try:
                    paddingLen = int(args[i+1], 10)
                except:
                    try:
                        paddingLen = int(args[i+1], 16)
                    except:
                        error("Error. '{}' is not a valid byte".format(args[i+1]))
                        return 
                if( paddingLen < 0 ):
                    error("Error. Padding length must be > 0")
                    return 
                i += 2
            elif( args[i] in [OPTION_OFFSET, OPTION_OFFSET_SHORT] ):
                if( seenOffset ):
                    error("Option '{}' should be used only once".format(args[i]))
                    return 
                if( i+1 >= len(args)):
                    error("Error. Missing offset after option '"+args[i]+"'")
                    return 
                seenOffset = True
                try:
                    offset = int(args[i+1], 10)
                except:
                    try:
                        offset = int(args[i+1], 16)
                    except:
                        error("Error. '{}' is not a valid offset".format(args[i+1]))
                        return 
                if( paddingLen < 0 ):
                    error("Error. Offset must be > 0")
                    return 
                i += 2
            elif( args[i] in [OPTION_HELP, OPTION_HELP_SHORT]):
                print_help()
                return 
            elif( args[i] in [OPTION_VERBOSE, OPTION_VERBOSE_SHORT]):
                seenVerbose = True
                i += 1
            else:
                error("Error. Unknown option '{}'".format(args[i]))
                return 
                
        # Parse a subcommand
        else:
            # Before parsing, check arguments 
            # Check arguments 
            if( seenPaddingByte and not seenPaddingLen ):
                error("Error. You have to specify padding length if you provide a padding byte. See the '-pl,--padding-length' option")
                return 
            # Parse and execute subcommand 
            verbose_mode(seenVerbose)
            payload = None
            
            ### set offset 
            if( not set_offset(offset)):
                error("Error. Your offset is too big :'(")
                return 
            try:
                ## Execute strategy
                if( args[i] == CMD_DELIVER_SHELLCODE ):
                    payload = dshell(args[i+1:], constraint, assertion, lmax=clmax)
                else:
                    error("Error. Unknown subcommand : '{}'".format(args[i]))
                    verbose_mode(False)
                    return 
            
                # Print result 
                verbose_mode(False)
                if( payload == "help"):
                    return 
                elif( not payload is None ):
                    print(string_bold("\n\tBuilt exploit - {} bytes\n".format(payload.len_bytes)))
                    # Normal output 
                    if( OUTPUT == OUTPUT_CONSOLE ):
                        payload_string = payload.strConsole(Arch.bits(), constraint.getBadBytes())
                    elif( OUTPUT == OUTPUT_PYTHON ):
                        payload_string = payload.strPython(Arch.bits(), constraint.getBadBytes(), \
                            noTab=seenOutfile, paddingByteStr=paddingByteStr, paddingLen=paddingLen)
                    # Output to file
                    if( outfile ):
                        try: 
                            f = open(outfile, "w")
                            f.write(remove_colors(payload_string))
                            f.close()
                            print(string_bold("\n\tWrote payload in file '{}'\n".format(outfile)))
                        except:
                            error("Error. Couldn't not write payload in file '{}'".format(outfile))
                    else:
                        print(payload_string)
                else:
                    error("Couldn't generate exploit")
            except Exception, e:
                reset_offset()
                raise sys.exc_info()[1], None, sys.exc_info()[2]
            
            ### reset offset 
            reset_offset()
            return 
                
    # Test parsing result 
    if( subcommand is None ):
        error("Missing subcommand")
        return




